import { TreeNode } from '@/domain/tree';
import { useEffect, useState } from 'react';

function _findSelectParent(node: TreeNode, selectKey: string, arr: Array<string> = []) {
    if (!node.children) {
        return selectKey === node.key;
    }
    let flag = false;
    node.children.forEach((child: any) => {
        if (_findSelectParent(child, selectKey, arr)) {
            flag = true;
        }
    });
    if (flag || node.isExpand) {
        arr.push(node.key);
    }
    return flag;
}

function _expandTreeNodes(nodes: Array<TreeNode>, arr: Array<string>): Array<TreeNode> {
    return nodes.map((node: any) => {
        if (!node.children) {
            return { ...node };
        }
        return { ...node, isExpand: arr.includes(node.key), children: _expandTreeNodes(node.children, arr) };
    });
}

/**
 * 树节点方法
 * @param {Object} initialData 树节点初始数据
 * @param {Array} defaultSelect 默认选中想
 * @param {Boolean} single 是否单选
 * @returns treeNodes: 处理后的树节点 handleNodeClick: 树节点点击函数
 */
const useTreeData = (initialData: Array<TreeNode> = [], defaultSelect: string) => {
    const [treeNodes, setTreeNodes] = useState(initialData);
    // 监听默认选中的改变
    useEffect(() => {
        let expandKeys: Array<string> = [];
        treeNodes.forEach((node) => {
            _findSelectParent(node, defaultSelect, expandKeys);
        });
        setTreeNodes(_expandTreeNodes(treeNodes, expandKeys));
    }, [defaultSelect]);
    // 处理节点的展开和折叠
    const updateTree = (nodes: Array<TreeNode>, key: string): Array<TreeNode> => {
        return nodes.map((n) => {
            if (n.key !== key && n.children) {
                return {
                    ...n,
                    children: updateTree(n.children, key),
                };
            }
            return { ...n, isExpand: n.key === key ? !n.isExpand : n.isExpand };
        });
    };
    return {
        treeNodes,
        setTreeNodes,
        updateTree,
    };
};

export default useTreeData;
